package com.yhkj.data.mq.handler;

import com.fasterxml.jackson.core.type.TypeReference;
import com.yhkj.data.datasource.JdbcManager;
import com.yhkj.data.mq.NativeMessageHandler;
import com.yhkj.data.mq.enums.TopicEnum;
import com.yhkj.data.mq.message.DataSourceConfig;
import com.yhkj.data.mq.message.SyncMessage;
import com.yhkj.data.mq.message.TableMetaData;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.collections4.ListUtils;
import org.apache.commons.collections4.keyvalue.MultiKey;
import org.apache.commons.collections4.map.MultiKeyMap;
import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyStatus;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @author <a href="mailto:idler41@163.con">linfuxin</a> created on 2023-03-06 16:07:35
 */
@Slf4j
@Component
public class SyncDataMessageHandler implements NativeMessageHandler<SyncMessage> {

    @Autowired
    private JdbcManager jdbcManager;

    @Value("${sync.pk.scan-batch-size:100}")
    private int scanPkBatchSize;

    private static final TypeReference<SyncMessage> DATA_TYPE = new TypeReference<SyncMessage>() {
    };

    @Override
    public TopicEnum getTopic() {
        return TopicEnum.GROUP_A_SYNC_DATA_GENERIC;
    }

    @Override
    public TypeReference<SyncMessage> getDataType() {
        return DATA_TYPE;
    }

    @Override
    public ConsumeConcurrentlyStatus process(SyncMessage content) {
        TableMetaData tableMetaData = content.getTableMetaData();
        DataSourceConfig dataSourceConfig = content.getDataSourceConfig();
        // 根据主键查询已存在数据
        final Integer[] keysIndex = tableMetaData.getKeysIndex();
        List<Object[]> existKeyList = Collections.emptyList();
        if (keysIndex != null) {
            List<List<Object>> keyDataList = Arrays.stream(content.getData()).map(i -> {
                List<Object> keyData = new ArrayList<>(keysIndex.length);
                for (Integer keyIdx : keysIndex) {
                    keyData.add(i[keyIdx]);
                }
                return keyData;
            }).collect(Collectors.toList());
            final String existPkSqlPrefix = buildExistKeySqlPrefix(tableMetaData.getTableName(), tableMetaData.getKeysColumnName());
            // 避免in查询参数过多，这里分批次查询
            existKeyList = doInBatch(keyDataList, scanPkBatchSize, (keyDataSubList) -> {
                // 参数由嵌套集合改为非嵌套方便设置参数
                List<Object> keyDataParamList = new ArrayList<>(keyDataSubList.size() * keysIndex.length);
                for (List<Object> keyDataRow : keyDataSubList) {
                    keyDataParamList.addAll(keyDataRow);
                }
                return jdbcManager.executeQueryPkList(dataSourceConfig, buildExistDataByKeySql(existPkSqlPrefix, keyDataSubList), tableMetaData.getKeysIndex().length, keyDataParamList);
            });
        }

        // insert、update 分组
        List<List<Object>> updateList = existKeyList.size() == 0 ? null : new ArrayList<>(existKeyList.size());
        List<List<Object>> insertList = new ArrayList<>(existKeyList.size() == 0 ? content.getData().length : content.getData().length - existKeyList.size());
        MultiKeyMap<Object, Object> multiKeyMap = toMultiKeyMap(existKeyList);
        for (Object[] rowData : content.getData()) {
            if (updateList != null && containsKey(rowData, multiKeyMap, keysIndex)) {
                List<Object> rowDataList = new ArrayList<>(rowData.length + 1);
                CollectionUtils.addAll(rowDataList, rowData);
                // 添加主键作为update参数
                for (Integer keyIndex : keysIndex) {
                    rowDataList.add(rowData[keyIndex]);
                }
                updateList.add(rowDataList);
            } else {
                List<Object> rowDataList = new ArrayList<>(rowData.length);
                CollectionUtils.addAll(rowDataList, rowData);
                insertList.add(rowDataList);
            }
        }

        // 批量执行sql
        int insertRowAffected = 0;
        int updateRowAffected = 0;
        String selectItem = String.join(",", tableMetaData.getColumnNames());
        if (CollectionUtils.isNotEmpty(insertList)) {
            String insertSql = buildInsertSql(tableMetaData.getTableName(), selectItem, tableMetaData.getColumnNames());
            insertRowAffected += jdbcManager.executeBatch(dataSourceConfig, insertSql, insertList);
        }
        if (CollectionUtils.isNotEmpty(updateList)) {
            String updateSql = buildUpdateSql(tableMetaData.getTableName(), tableMetaData.getColumnNames(), tableMetaData.getKeysColumnName());
            updateRowAffected += jdbcManager.executeBatch(dataSourceConfig, updateSql, updateList);
        }
        log.info("{}表同步完成!新增{}条，更新{}条", tableMetaData.getTableName(), insertRowAffected, updateRowAffected);
        return ConsumeConcurrentlyStatus.CONSUME_SUCCESS;
    }

    private <T, R> List<R> doInBatch(List<T> data, int batchSize, Function<List<T>, List<R>> function) {
        if (data.size() < batchSize) {
            return function.apply(data);
        }
        List<List<T>> partList = ListUtils.partition(data, batchSize);
        List<R> result = new LinkedList<>();
        for (List<T> subList : partList) {
            result.addAll(function.apply(subList));
        }
        return result;
    }

    private MultiKeyMap<Object, Object> toMultiKeyMap(List<Object[]> existKeyList) {
        MultiKeyMap<Object, Object> multiKeyMap = new MultiKeyMap<>();
        for (Object[] existKeys : existKeyList) {
            multiKeyMap.put(new MultiKey<>(existKeys), 1);
        }
        return multiKeyMap;
    }

    private boolean containsKey(Object[] rowData, MultiKeyMap<Object, Object> multiKeyMap, Integer[] keyIndex) {
        int keySize = keyIndex.length;
        Object[] keyData = new Object[keySize];
        for (int i = 0; i < keySize; i++) {
            keyData[i] = rowData[keyIndex[i]];
        }
        return multiKeyMap.containsKey(new MultiKey<>(keyData));
    }

    private String buildExistKeySqlPrefix(String tableName, String[] keyColumns) {
        return "SELECT " + String.join(",", keyColumns) + " FROM " + tableName + " WHERE ("
                + String.join(",", keyColumns) + ") IN ( ";
    }

    private String buildExistPkSql(String tableName, String[] keyColumns, List<List<Object>> keyDataList) {
        return "SELECT " + String.join(",", keyColumns) + " FROM " + tableName + " WHERE ("
                + String.join(",", keyColumns) + ") IN ( " + keyDataList.stream()
                .map(i -> "(" + i.stream().map(j -> "?").collect(Collectors.joining(",")) + ")")
                .collect(Collectors.joining(",")) + " )";
    }

    private String buildExistDataByKeySql(String existKeySqlPrefix, List<List<Object>> keyDataList) {
        return existKeySqlPrefix + keyDataList.stream()
                .map(i -> "(" + i.stream().map(j -> "?").collect(Collectors.joining(",")) + ")")
                .collect(Collectors.joining(",")) + " )";
    }

    private String buildInsertSql(String tableName, String selectItems, String[] columnNames) {
        return "INSERT INTO " + tableName + " ( " + selectItems + " ) VALUES  ( " +
                Arrays.stream(columnNames)
                        .map(i -> "?").collect(Collectors.joining(",")) + " ) ";
    }

    private String buildUpdateSql(String tableName, String[] columnNames, String[] keyColumnNames) {
        return "UPDATE " + tableName + " SET " +
                Arrays.stream(columnNames).map(i -> i + "=?").collect(Collectors.joining(",")) +
                " WHERE (" + String.join(",", keyColumnNames)
                + ") = (" + Arrays.stream(keyColumnNames).map(i -> "?").collect(Collectors.joining(",")) + ")";
    }
}
